""" Full assembly of the parts to form the complete network """

import torch.nn.functional as F
import torch.nn as nn 
import torch 


# from .unet_parts import *

# class Discriminator(nn.Module):
#     def __init__(self, n_channels):
#         super(Discriminator, self).__init__()
#         self.n_channels = n_channels
#         input_in_channels = 2 * n_channels

#         # self.inc = DoubleConv(n_channels, 64)
#         self.down1 = Down(input_in_channels, 64)
#         self.down2 = Down(64, 128)
#         self.down3 = Down(128, 256)
#         self.outc = OutConv(256, 1)

#     def forward(self, x1, x2):
#         x = torch.cat((x1, x2), dim=1)
#         # print(x1.shape)
#         x1 = self.down1(x)# (128)
#         # print(x2.shape)
#         x2 = self.down2(x1) # (256)
#         # print(x3.shape)
#         x3 = self.down3(x2) # (512) 
#         # print(x4.shape)

#         logits = self.outc(x3)
#         # print(logits.shape)
#         return logits


class Discriminator(nn.Module):
    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
        super(Discriminator, self).__init__()
        kw = 4
        padw = 1
        sequence = [nn.Conv2d(2*input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, n_layers):  # gradually increase the number of filters
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** n, 8)
            sequence += [
                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw),
                norm_layer(ndf * nf_mult),
                nn.LeakyReLU(0.2, True)
            ]

        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8)
        sequence += [
            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw),
            norm_layer(ndf * nf_mult),
            nn.LeakyReLU(0.2, True)
        ]

        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
        self.model = nn.Sequential(*sequence)

    def forward(self, x1, x2):
        """Standard forward."""
        x = torch.cat((x1, x2), dim=1)
        return self.model(x)


if __name__ == "__main__":
    
    t1 = torch.rand((1, 3, 200, 200))
    t2 = torch.rand((1, 3, 200, 200))

    dis = Discriminator(3)
    out = dis(t1, t1)

    print(out.shape)